import json
import epics
import os
from typing import Dict, List, Union
import argparse
import sys

class ParameterSetter:
    def __init__(self, config_file: str = "config.json", default_params_file: str = "default_parameters.json"):
        """初始化参数设置器
        
        Args:
            config_file: 配置文件路径
            default_params_file: 默认参数文件路径
        """
        self.config_file = config_file
        self.default_params_file = default_params_file
        self.config = self._load_config()
        self.default_params = self._load_default_params()
        self._generate_cavity_configs()
        
    def _load_config(self) -> dict:
        """加载配置文件"""
        if not os.path.exists(self.config_file):
            raise FileNotFoundError(f"配置文件 {self.config_file} 不存在")
        
        with open(self.config_file, 'r') as f:
            return json.load(f)
    
    def _load_default_params(self) -> dict:
        """加载默认参数文件"""
        if not os.path.exists(self.default_params_file):
            print(f"警告: 默认参数文件 {self.default_params_file} 不存在，将使用空字典")
            return {}
        
        with open(self.default_params_file, 'r') as f:
            return json.load(f)
    
    def _save_config(self):
        """保存配置到文件"""
        with open(self.config_file, 'w') as f:
            json.dump(self.config, f, indent=4)
    
    def _save_default_params(self):
        """保存默认参数到文件"""
        with open(self.default_params_file, 'w') as f:
            json.dump(self.default_params, f, indent=4)
    
    def _convert_alias_to_pv(self, cavity_name: str) -> str:
        """将别名格式转换为PV格式
        
        Args:
            cavity_name: 别名格式的腔体名称，如 "CM1-1"
            
        Returns:
            PV格式的腔体名称，如 "CM01-01"
        """
        if not cavity_name.startswith("CM"):
            return cavity_name
        parts = cavity_name.split("-")
        if len(parts) != 2:
            return cavity_name
        module = parts[0][2:].zfill(2)  # 提取模块号并补零
        unit = parts[1].zfill(2)  # 单元号补零
        return f"CM{module}-{unit}"
    
    def _convert_pv_to_alias(self, cavity_name: str) -> str:
        """将PV格式转换为别名格式
        
        Args:
            cavity_name: PV格式的腔体名称，如 "CM01-01"
            
        Returns:
            别名格式的腔体名称，如 "CM1-1"
        """
        if not cavity_name.startswith("CM"):
            return cavity_name
        parts = cavity_name.split("-")
        if len(parts) != 2:
            return cavity_name
        module = str(int(parts[0][2:]))  # 提取模块号并去除前导零
        unit = str(int(parts[1]))  # 单元号去除前导零
        return f"CM{module}-{unit}"
    
    def _generate_cavity_configs(self):
        """根据配置模式生成所有超导腔的PV配置"""
        self.cavity_configs = {}
        pattern = self.config["PV_PATTERN"]
        
        for module in self.config["MODULES"]:
            for unit in self.config["UNITS"]:
                cavity_name = f"CM{module}-{unit}"
                pv_config = {}
                
                for param in self.config["PARAMETERS"]:
                    pv_name = (pattern["prefix"].format(module=module, unit=unit) + 
                             pattern["suffix"].format(param=param))
                    pv_config[f"offset{param}"] = pv_name
                
                self.cavity_configs[cavity_name] = pv_config
    
    def get_cavity_pvs(self, cavity_name: str) -> Dict[str, str]:
        """获取指定超导腔的所有PV名称
        
        Args:
            cavity_name: 超导腔名称，如 "CM1-1" 或 "CM01-01"
            
        Returns:
            包含PV名称的字典
        """
        pv_name = self._convert_alias_to_pv(cavity_name)
        if pv_name not in self.cavity_configs:
            raise ValueError(f"超导腔 {cavity_name} 不存在")
        return self.cavity_configs[pv_name]
    
    def get_module_cavities(self, module: str) -> List[str]:
        """获取指定恒温器的所有超导腔名称
        
        Args:
            module: 恒温器编号，如 "1"
            
        Returns:
            该恒温器下所有超导腔的名称列表（使用别名格式）
        """
        module = module.zfill(2)  # 确保模块编号是两位数
        pv_cavities = [cavity for cavity in self.cavity_configs.keys() 
                      if cavity.startswith(f"CM{module}-")]
        return [self._convert_pv_to_alias(cavity) for cavity in pv_cavities]
    
    def get_default_parameters(self, cavity_name: str) -> Dict[str, float]:
        """获取指定超导腔的默认参数
        
        Args:
            cavity_name: 超导腔名称（使用别名格式）
            
        Returns:
            包含默认参数的字典
        """
        return self.default_params.get(cavity_name, {})
    
    def set_default_parameters(self, cavity_name: str, parameters: Dict[str, float]):
        """设置指定超导腔的默认参数
        
        Args:
            cavity_name: 超导腔名称（使用别名格式）
            parameters: 参数字典
        """
        self.default_params[cavity_name] = parameters
        self._save_default_params()
    
    def set_module_default_parameters(self, module: str, parameters: Dict[str, float]):
        """设置指定恒温器所有超导腔的默认参数
        
        Args:
            module: 恒温器编号，如 "1"
            parameters: 参数字典
        """
        cavities = self.get_module_cavities(module)
        for cavity in cavities:
            self.set_default_parameters(cavity, parameters)
    
    def set_pv_value(self, pv_name: str, value: float, timeout: float = 5.0) -> bool:
        """设置单个PV的值
        
        Args:
            pv_name: PV名称
            value: 要设置的值
            timeout: 超时时间（秒）
            
        Returns:
            是否设置成功
        """
        try:
            pv = epics.PV(pv_name)
            if not pv.wait_for_connection(timeout=timeout):
                print(f"无法连接到PV: {pv_name}")
                return False
            
            pv.put(value)
            return True
        except Exception as e:
            print(f"设置PV {pv_name} 时出错: {str(e)}")
            return False
    
    def set_cavity_parameters(self, cavity_name: str, parameters: Dict[str, float] = None) -> Dict[str, dict]:
        """设置指定超导腔的多个参数
        
        Args:
            cavity_name: 超导腔名称（使用别名格式）
            parameters: 参数字典，键为参数名（如"offset4"），值为要设置的值
                       如果为None，则使用默认参数
        
        Returns:
            每个参数设置的详细结果字典
        """
        pv_name = self._convert_alias_to_pv(cavity_name)
        pvs = self.get_cavity_pvs(pv_name)
        results = {}
        
        # 如果没有提供参数，使用默认参数
        if parameters is None:
            parameters = self.get_default_parameters(cavity_name)
            if not parameters:
                print(f"警告: 超导腔 {cavity_name} 没有默认参数")
                return {}
        
        for param_name, value in parameters.items():
            if param_name not in pvs:
                print(f"参数 {param_name} 不存在于超导腔 {cavity_name}")
                results[param_name] = {"success": False, "pv": None, "value": value}
                continue
                
            pv_name = pvs[param_name]
            success = self.set_pv_value(pv_name, value)
            results[param_name] = {"success": success, "pv": pv_name, "value": value}
        
        return results
    
    def set_module_parameters(self, module: str, parameters: Dict[str, float] = None) -> Dict[str, Dict[str, dict]]:
        """设置指定恒温器所有超导腔的参数
        
        Args:
            module: 恒温器编号，如 "1"
            parameters: 参数字典，如果为None则使用默认参数
            
        Returns:
            每个超导腔的设置结果字典
        """
        cavities = self.get_module_cavities(module)
        results = {}
        
        for cavity in cavities:
            results[cavity] = self.set_cavity_parameters(cavity, parameters)
            
        return results
    
    def get_available_cavities(self) -> List[str]:
        """获取所有可用的超导腔名称（使用别名格式）"""
        return [self._convert_pv_to_alias(cavity) for cavity in self.cavity_configs.keys()]
    
    def add_module(self, module: str):
        """添加新的恒温器
        
        Args:
            module: 恒温器编号（如 "5"）
        """
        module = module.zfill(2)  # 确保模块编号是两位数
        if module not in self.config["MODULES"]:
            self.config["MODULES"].append(module)
            self._save_config()
            self._generate_cavity_configs()
    
    def add_unit(self, unit: str):
        """添加新的单元
        
        Args:
            unit: 单元编号（如 "6"）
        """
        unit = unit.zfill(2)  # 确保单元编号是两位数
        if unit not in self.config["UNITS"]:
            self.config["UNITS"].append(unit)
            self._save_config()
            self._generate_cavity_configs()
    
    def add_parameter(self, param: str):
        """添加新的参数
        
        Args:
            param: 参数编号（如 "9"）
        """
        if param not in self.config["PARAMETERS"]:
            self.config["PARAMETERS"].append(param)
            self._save_config()
            self._generate_cavity_configs()

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='超导腔参数设置工具')
    subparsers = parser.add_subparsers(dest='command', help='可用命令')
    
    # 列出所有可用的超导腔
    list_parser = subparsers.add_parser('list', help='列出所有可用的超导腔')
    
    # 设置单个超导腔的参数
    set_cavity_parser = subparsers.add_parser('set-cavity', help='设置单个超导腔的参数')
    set_cavity_parser.add_argument('cavity', help='超导腔名称，如 CM1-1')
    set_cavity_parser.add_argument('--params', nargs='+', help='参数设置，格式为 offset4=16.0 offset5=26.0（可选，默认使用default_parameters.json中的参数）')
    
    # 设置整个模块的参数
    set_module_parser = subparsers.add_parser('set-module', help='设置整个模块的参数')
    set_module_parser.add_argument('module', help='模块编号，如 1')
    set_module_parser.add_argument('--params', nargs='+', help='参数设置，格式为 offset4=16.0 offset5=26.0（可选，默认使用default_parameters.json中的参数）')
    
    return parser.parse_args()

def parse_parameters(params_list):
    """解析参数字符串列表为字典"""
    if not params_list:
        return None
    params = {}
    for param_str in params_list:
        try:
            key, value = param_str.split('=')
            params[key] = float(value)
        except ValueError:
            print(f"警告: 无法解析参数 {param_str}，已跳过")
    return params

def main():
    """主函数"""
    args = parse_args()
    setter = ParameterSetter()
    
    if args.command == 'list':
        cavities = setter.get_available_cavities()
        print("可用的超导腔:")
        for cavity in cavities:
            print(f"  {cavity}")
    
    elif args.command == 'set-cavity':
        # 如果没有提供参数，使用默认参数
        if not args.params:
            print(f"使用 default_parameters.json 中的参数设置超导腔 {args.cavity}")
            results = setter.set_cavity_parameters(args.cavity)
        else:
            params = parse_parameters(args.params)
            results = setter.set_cavity_parameters(args.cavity, params)
        
        print(f"设置超导腔 {args.cavity} 的结果:")
        for param, info in results.items():
            if info["success"]:
                print(f"  {param}: 成功 PV: {info['pv']}，值: {info['value']}")
            else:
                print(f"  {param}: 失败")
    
    elif args.command == 'set-module':
        # 如果没有提供参数，使用默认参数
        if not args.params:
            print(f"使用 default_parameters.json 中的参数设置模块 {args.module}")
            results = setter.set_module_parameters(args.module)
        else:
            params = parse_parameters(args.params)
            results = setter.set_module_parameters(args.module, params)
        
        print(f"设置模块 {args.module} 的结果:")
        for cavity, cavity_results in results.items():
            print(f"\n超导腔 {cavity}:")
            for param, info in cavity_results.items():
                if info["success"]:
                    print(f"  {param}: 成功 PV: {info['pv']}，值: {info['value']}")
                else:
                    print(f"  {param}: 失败")
    
    else:
        print("请指定要执行的命令。使用 -h 查看帮助信息。")
        sys.exit(1)

if __name__ == "__main__":
    main() 